import numpy as np
import yaml

# 生成开阔室内场景的PGM地图
def generate_open_space_pgm(output_pgm="open_space_map.pgm"):
    # 地图尺寸：1000（宽）× 800（高）像素 → 更开阔的空间
    width, height = 1000, 800
    # 初始化全白画布（自由空间，灰度值255）
    map_data = np.ones((height, width), dtype=np.uint8) * 255

    # ==================== 1. 绘制外围封闭墙体（黑色，灰度0） ====================
    wall_thickness = 20
    map_data[:wall_thickness, :] = 0        # 顶部墙体
    map_data[-wall_thickness:, :] = 0       # 底部墙体
    map_data[:, :wall_thickness] = 0        # 左侧墙体
    map_data[:, -wall_thickness:] = 0       # 右侧墙体

    # ==================== 2. 绘制水平直线障碍物（细长，宽度2像素） ====================
    line_thickness = 2
    # 上方水平直线
    map_data[200:200+line_thickness, 300:700] = 0
    # 下方水平直线
    map_data[600:600+line_thickness, 300:700] = 0

    # ==================== 3. 绘制垂直直线障碍物（细长，宽度2像素） ====================
    # 左侧垂直直线
    map_data[250:550, 250:250+line_thickness] = 0
    # 右侧垂直直线
    map_data[250:550, 750:750+line_thickness] = 0

    # ==================== 4. 绘制散点小矩形障碍物（15x15像素，体积小） ====================
    obs_size = 15  # 小矩形尺寸
    # 顶部散点
    scatter_obs = [
        (150, 150), (400, 120), (600, 150), (850, 180),
        # 中间散点
        (350, 350), (500, 450), (650, 380),
        # 底部散点
        (200, 650), (450, 680), (700, 650), (800, 750)
    ]
    for x, y in scatter_obs:
        map_data[y:y+obs_size, x:x+obs_size] = 0

    # ==================== 5. 绘制四周小矩形组（模拟墙边物体） ====================
    # 左上角小矩形
    map_data[50:80, 50:100] = 0
    # 右上角小矩形
    map_data[50:80, 900:950] = 0
    # 左下角小矩形
    map_data[700:750, 50:100] = 0
    # 右下角小矩形
    map_data[700:750, 900:950] = 0

    # 保存为PGM文件（二进制格式）
    with open(output_pgm, 'wb') as f:
        f.write(b'P5\n# Open space indoor map with small obstacles\n')
        f.write(f'{width} {height}\n255\n'.encode('ascii'))
        f.write(map_data.tobytes())
    print(f" 开阔场景PGM地图已生成：{output_pgm}")

# 生成配套的YAML配置文件
def generate_open_space_yaml(pgm_path="open_space_map.pgm", output_yaml="open_space_map.yaml"):
    config = {
        "image": pgm_path,            # PGM文件相对路径
        "resolution": 0.1,            # 分辨率：0.1米/像素 → 地图实际100m×80m
        "origin": [-50.0, -40.0, 0.0],# 地图原点（左下角世界坐标：x=-50m, y=-40m）
        "negate": 0,                  # 不反转（黑色=障碍，白色=自由）
        "occupied_thresh": 0.65,      # 障碍灰度阈值（≤89视为障碍）
        "free_thresh": 0.25,          # 自由灰度阈值（≥64视为自由）
        "mode": "trinary",            # 三值模式（自由/障碍/未知）
        "dilate_kernel_size": 3       # 障碍膨胀核大小（路径规划用）
    }
    with open(output_yaml, "w") as f:
        yaml.dump(config, f, sort_keys=False)  # 保持字段顺序易读
    print(f" 配套YAML文件已生成：{output_yaml}")

if __name__ == "__main__":
    generate_open_space_pgm()
    generate_open_space_yaml()